"""
    EnsembleInflation

Comprehensive ensemble inflation system for GSI EnKF data assimilation.
This module implements various inflation techniques to maintain ensemble spread
and prevent filter divergence. Based on the GSI/EnKF Fortran implementation
from inflation.f90.

Mathematical Framework:
The inflation system uses multiplicative inflation where the inflation factor
is computed as:
```
r = analpertwt × ((σ_prior - σ_posterior) / σ_posterior) + 1
```

Where:
- σ_prior: Prior ensemble standard deviation
- σ_posterior: Posterior ensemble standard deviation (before inflation)
- analpertwt: Regional inflation weight (NH/SH/TR specific)
- r: Inflation factor applied to ensemble perturbations

Inflation Methods:
1. **Multiplicative Inflation**: Standard spread-based inflation
2. **Relaxation-to-Prior (RTP)**: Blends analysis with prior ensemble
3. **Adaptive Inflation**: Uses innovation statistics for automatic tuning
4. **Regional Inflation**: Different coefficients for NH/SH/TR regions
5. **Spectral Smoothing**: Gaussian filter for inflation field smoothing
"""
module EnsembleInflation

using LinearAlgebra
using Statistics
using FFTW
using Random
using ..GSICoreAnalysis: AbstractAnalysisConfig, AbstractStateVector
using ..StateVectors: StateVector
using ..GridOperations: GridConfiguration, spectral_transform, inverse_spectral_transform

export InflationConfiguration, InflationResult, InflationDiagnostics
export inflate_ensemble!, compute_inflation_factor, adaptive_inflation!
export regional_inflation_weights, smooth_inflation_field
export relaxation_to_prior!, inflation_diagnostics

# Regional area definitions
const REGION_NH = 1  # Northern Hemisphere
const REGION_SH = 2  # Southern Hemisphere  
const REGION_TR = 3  # Tropics

"""
    InflationConfiguration{T<:AbstractFloat}

Configuration parameters for ensemble inflation system.
"""
struct InflationConfiguration{T<:AbstractFloat}
    # Regional inflation weights
    analpertwt_nh::T          # Northern hemisphere inflation weight
    analpertwt_sh::T          # Southern hemisphere inflation weight  
    analpertwt_tr::T          # Tropical inflation weight
    
    # Inflation bounds
    covinflatemin::T          # Minimum allowed inflation factor (default: 0.0)
    covinflatemax::T          # Maximum allowed inflation factor (default: 10.0)
    
    # Smoothing parameters
    smoothparm::T             # Gaussian smoothing e-folding scale (grid units)
    enable_smoothing::Bool    # Whether to apply spectral smoothing
    
    # Regional boundaries (degrees)
    latbound_nh::T            # Northern hemisphere boundary (default: 20.0)
    latbound_sh::T            # Southern hemisphere boundary (default: -20.0)
    
    # Adaptive inflation parameters
    adaptive_inflation::Bool   # Enable adaptive inflation
    innovation_threshold::T    # Innovation-based adaptation threshold
    update_frequency::Int      # How often to update inflation (analysis cycles)
    
    # Relaxation-to-prior parameters
    rtp_inflation::Bool       # Use relaxation-to-prior method
    rtp_alpha::T             # RTP mixing parameter (0-1)
    
    # Diagnostic output
    save_inflation::Bool      # Save inflation fields to file
    compute_diagnostics::Bool # Compute area-mean statistics
end

"""
    InflationResult{T<:AbstractFloat}

Results from ensemble inflation process.
"""
struct InflationResult{T<:AbstractFloat}
    inflation_factors::Array{T,3}    # Inflation factors (nx,ny,nlevels)
    inflated_perturbations::Vector{StateVector{T}}  # Inflated ensemble perturbations
    diagnostics::Union{InflationDiagnostics{T}, Nothing}
end

"""
    InflationDiagnostics{T<:AbstractFloat}

Diagnostic information from inflation process.
"""
struct InflationDiagnostics{T<:AbstractFloat}
    # Regional statistics
    mean_inflation_nh::T      # Mean inflation factor - Northern Hemisphere
    mean_inflation_sh::T      # Mean inflation factor - Southern Hemisphere
    mean_inflation_tr::T      # Mean inflation factor - Tropics
    
    # Global statistics
    min_inflation::T          # Minimum inflation factor globally
    max_inflation::T          # Maximum inflation factor globally
    mean_inflation::T         # Global mean inflation factor
    
    # Spread statistics
    prior_spread_mean::T      # Mean prior ensemble spread
    posterior_spread_mean::T  # Mean posterior ensemble spread
    final_spread_mean::T      # Mean final (inflated) ensemble spread
    
    # Regional area weights (for proper averaging)
    area_weights_nh::T        # Area weight for NH
    area_weights_sh::T        # Area weight for SH
    area_weights_tr::T        # Area weight for TR
end

"""
    inflate_ensemble!(ensemble_perturbations, prior_perturbations, config, grid_config, latitudes)

Apply multiplicative inflation to ensemble perturbations.

# Arguments
- `ensemble_perturbations`: Posterior ensemble perturbations (modified in-place)
- `prior_perturbations`: Prior ensemble perturbations (for comparison)
- `config`: InflationConfiguration
- `grid_config`: GridConfiguration for spatial operations
- `latitudes`: Latitude values for regional weighting

# Returns
- `InflationResult`: Complete inflation results and diagnostics
"""
function inflate_ensemble!(
    ensemble_perturbations::Vector{StateVector{T}},
    prior_perturbations::Vector{StateVector{T}},
    config::InflationConfiguration{T},
    grid_config::GridConfiguration{T},
    latitudes::AbstractArray{T}
) where T<:AbstractFloat
    
    # Check if inflation is requested
    if abs(config.analpertwt_nh) < 1e-5 && abs(config.analpertwt_sh) < 1e-5 && abs(config.analpertwt_tr) < 1e-5
        return InflationResult(
            zeros(T, grid_config.nx, grid_config.ny, grid_config.nlevels),
            ensemble_perturbations,
            nothing
        )
    end
    
    nensmembers = length(ensemble_perturbations)
    nx, ny, nlevels = grid_config.nx, grid_config.ny, grid_config.nlevels
    
    # Compute ensemble statistics
    prior_mean = compute_ensemble_mean(prior_perturbations)
    posterior_mean = compute_ensemble_mean(ensemble_perturbations)
    
    # Compute standard deviations
    prior_std = compute_ensemble_std(prior_perturbations, prior_mean)
    posterior_std = compute_ensemble_std(ensemble_perturbations, posterior_mean)
    
    # Initialize inflation factors
    inflation_factors = ones(T, nx, ny, nlevels)
    
    # Compute regional weights
    regional_weights = compute_regional_weights(latitudes, config)
    
    # Apply inflation method
    if config.rtp_inflation
        apply_relaxation_to_prior!(ensemble_perturbations, prior_perturbations, 
                                 posterior_mean, config, regional_weights)
    else
        compute_multiplicative_inflation!(inflation_factors, prior_std, posterior_std,
                                        config, regional_weights)
        
        # Apply spectral smoothing if requested
        if config.enable_smoothing && config.smoothparm > 0
            smooth_inflation_field!(inflation_factors, config.smoothparm, grid_config)
        end
        
        # Apply inflation bounds
        clamp!(inflation_factors, config.covinflatemin, config.covinflatemax)
        
        # Apply inflation to perturbations
        apply_inflation_factors!(ensemble_perturbations, inflation_factors, posterior_mean)
    end
    
    # Compute diagnostics if requested
    diagnostics = nothing
    if config.compute_diagnostics
        final_std = compute_ensemble_std(ensemble_perturbations, posterior_mean)
        diagnostics = compute_inflation_diagnostics(inflation_factors, prior_std, 
                                                  posterior_std, final_std, 
                                                  regional_weights, latitudes, config)
    end
    
    return InflationResult(inflation_factors, ensemble_perturbations, diagnostics)
end

"""
    compute_multiplicative_inflation!(inflation_factors, prior_std, posterior_std, config, weights)

Compute multiplicative inflation factors based on spread reduction.
"""
function compute_multiplicative_inflation!(
    inflation_factors::Array{T,3},
    prior_std::Array{T,3},
    posterior_std::Array{T,3},
    config::InflationConfiguration{T},
    regional_weights::Array{T,3}
) where T<:AbstractFloat
    
    nx, ny, nlevels = size(inflation_factors)
    
    @inbounds for k in 1:nlevels
        for j in 1:ny
            for i in 1:nx
                if posterior_std[i,j,k] > 1e-10  # Avoid division by zero
                    # Get regional inflation weight
                    weight = regional_weights[i,j,k]
                    
                    # Compute inflation factor
                    spread_reduction = (prior_std[i,j,k] - posterior_std[i,j,k]) / posterior_std[i,j,k]
                    inflation_factors[i,j,k] = weight * spread_reduction + one(T)
                else
                    inflation_factors[i,j,k] = one(T)
                end
            end
        end
    end
end

"""
    apply_relaxation_to_prior!(ensemble_perturbations, prior_perturbations, mean, config, weights)

Apply relaxation-to-prior inflation method.
"""
function apply_relaxation_to_prior!(
    ensemble_perturbations::Vector{StateVector{T}},
    prior_perturbations::Vector{StateVector{T}},
    posterior_mean::StateVector{T},
    config::InflationConfiguration{T},
    regional_weights::Array{T,3}
) where T<:AbstractFloat
    
    nensmembers = length(ensemble_perturbations)
    alpha = abs(config.rtp_alpha)  # Mixing parameter
    
    for imem in 1:nensmembers
        posterior_pert = ensemble_perturbations[imem] - posterior_mean
        prior_pert = prior_perturbations[imem] - posterior_mean
        
        # Apply regional weighting for RTP
        for field_name in fieldnames(typeof(posterior_pert))
            posterior_field = getfield(posterior_pert, field_name)
            prior_field = getfield(prior_pert, field_name)
            
            if isa(posterior_field, AbstractArray)
                @. posterior_field = (1 - alpha) * posterior_field + alpha * prior_field
            end
        end
        
        # Update ensemble member
        ensemble_perturbations[imem] = posterior_mean + posterior_pert
    end
end

"""
    compute_regional_weights(latitudes, config)

Compute regional inflation weights based on latitude.
"""
function compute_regional_weights(
    latitudes::AbstractArray{T},
    config::InflationConfiguration{T}
) where T<:AbstractFloat
    
    nx, ny = size(latitudes)
    nlevels = 1  # Will be broadcasted
    weights = zeros(T, nx, ny, nlevels)
    
    @inbounds for j in 1:ny
        for i in 1:nx
            lat = latitudes[i,j]
            
            if lat >= config.latbound_nh
                # Northern Hemisphere
                weights[i,j,1] = config.analpertwt_nh
            elseif lat <= config.latbound_sh
                # Southern Hemisphere
                weights[i,j,1] = config.analpertwt_sh
            else
                # Tropics
                weights[i,j,1] = config.analpertwt_tr
            end
        end
    end
    
    # Broadcast to all levels
    return repeat(weights, 1, 1, nlevels)
end

"""
    smooth_inflation_field!(inflation_factors, smoothparm, grid_config)

Apply Gaussian spectral smoothing to inflation field.
"""
function smooth_inflation_field!(
    inflation_factors::Array{T,3},
    smoothparm::T,
    grid_config::GridConfiguration{T}
) where T<:AbstractFloat
    
    nx, ny, nlevels = size(inflation_factors)
    
    # Apply smoothing to each level
    for k in 1:nlevels
        field = view(inflation_factors, :, :, k)
        
        # Convert to spectral space
        field_spec = spectral_transform(field, grid_config)
        
        # Apply Gaussian filter
        apply_gaussian_filter!(field_spec, smoothparm, grid_config)
        
        # Convert back to grid space
        field_smooth = inverse_spectral_transform(field_spec, grid_config)
        
        # Update field
        inflation_factors[:, :, k] = field_smooth
    end
end

"""
    apply_gaussian_filter!(field_spectral, smoothparm, grid_config)

Apply Gaussian smoothing filter in spectral space.
"""
function apply_gaussian_filter!(
    field_spectral::Array{Complex{T},2},
    smoothparm::T,
    grid_config::GridConfiguration{T}
) where T<:AbstractFloat
    
    nlon_spec, nlat_spec = size(field_spectral)
    
    # Compute smoothing weights
    for n in 1:nlat_spec
        for m in 1:nlon_spec
            # Total wavenumber
            k_total = sqrt(T(m-1)^2 + T(n-1)^2)
            
            # Gaussian smoothing factor
            smooth_factor = exp(-(k_total * smoothparm)^2 / 4)
            
            field_spectral[m,n] *= smooth_factor
        end
    end
end

"""
    apply_inflation_factors!(ensemble_perturbations, inflation_factors, mean)

Apply computed inflation factors to ensemble perturbations.
"""
function apply_inflation_factors!(
    ensemble_perturbations::Vector{StateVector{T}},
    inflation_factors::Array{T,3},
    ensemble_mean::StateVector{T}
) where T<:AbstractFloat
    
    nensmembers = length(ensemble_perturbations)
    
    for imem in 1:nensmembers
        # Compute perturbation
        perturbation = ensemble_perturbations[imem] - ensemble_mean
        
        # Apply inflation to each field
        for field_name in fieldnames(typeof(perturbation))
            field = getfield(perturbation, field_name)
            
            if isa(field, Array{T,3})
                # 3D field - apply inflation factors directly
                field .*= inflation_factors
            elseif isa(field, Array{T,2})
                # 2D field - use surface level inflation
                field .*= inflation_factors[:,:,1]
            end
        end
        
        # Update ensemble member
        ensemble_perturbations[imem] = ensemble_mean + perturbation
    end
end

"""
    compute_ensemble_mean(ensemble_perturbations)

Compute ensemble mean state vector.
"""
function compute_ensemble_mean(ensemble_perturbations::Vector{StateVector{T}}) where T<:AbstractFloat
    nensmembers = length(ensemble_perturbations)
    mean_state = deepcopy(ensemble_perturbations[1])
    
    # Initialize to zero
    for field_name in fieldnames(typeof(mean_state))
        field = getfield(mean_state, field_name)
        if isa(field, AbstractArray)
            fill!(field, zero(T))
        end
    end
    
    # Sum all ensemble members
    for imem in 1:nensmembers
        for field_name in fieldnames(typeof(mean_state))
            mean_field = getfield(mean_state, field_name)
            ens_field = getfield(ensemble_perturbations[imem], field_name)
            
            if isa(mean_field, AbstractArray)
                mean_field .+= ens_field
            end
        end
    end
    
    # Divide by ensemble size
    inv_nensmembers = one(T) / T(nensmembers)
    for field_name in fieldnames(typeof(mean_state))
        field = getfield(mean_state, field_name)
        if isa(field, AbstractArray)
            field .*= inv_nensmembers
        end
    end
    
    return mean_state
end

"""
    compute_ensemble_std(ensemble_perturbations, ensemble_mean)

Compute ensemble standard deviation.
"""
function compute_ensemble_std(
    ensemble_perturbations::Vector{StateVector{T}},
    ensemble_mean::StateVector{T}
) where T<:AbstractFloat
    
    nensmembers = length(ensemble_perturbations)
    inv_nensmembers_minus_one = one(T) / T(nensmembers - 1)
    
    # Initialize standard deviation arrays (assuming 3D primary fields)
    # This is simplified - in practice would need to handle all state vector fields
    first_member = ensemble_perturbations[1]
    
    # Find a representative 3D field to determine dimensions
    field_3d = nothing
    for field_name in fieldnames(typeof(first_member))
        field = getfield(first_member, field_name)
        if isa(field, Array{T,3})
            field_3d = field
            break
        end
    end
    
    if field_3d === nothing
        error("No 3D fields found in state vector for std computation")
    end
    
    std_field = zeros(T, size(field_3d))
    
    # Compute variance
    for imem in 1:nensmembers
        member = ensemble_perturbations[imem]
        mean_member = ensemble_mean
        
        for field_name in fieldnames(typeof(member))
            if isa(getfield(member, field_name), Array{T,3})
                diff = getfield(member, field_name) - getfield(mean_member, field_name)
                std_field .+= diff .* diff
                break  # Just use first 3D field for simplicity
            end
        end
    end
    
    # Convert to standard deviation
    std_field .*= inv_nensmembers_minus_one
    std_field .= sqrt.(std_field)
    
    return std_field
end

"""
    compute_inflation_diagnostics(inflation_factors, prior_std, posterior_std, final_std, weights, latitudes, config)

Compute comprehensive inflation diagnostics.
"""
function compute_inflation_diagnostics(
    inflation_factors::Array{T,3},
    prior_std::Array{T,3},
    posterior_std::Array{T,3},
    final_std::Array{T,3},
    regional_weights::Array{T,3},
    latitudes::AbstractArray{T},
    config::InflationConfiguration{T}
) where T<:AbstractFloat
    
    nx, ny, nlevels = size(inflation_factors)
    
    # Compute area weights (cosine of latitude)
    area_weights = zeros(T, nx, ny)
    for j in 1:ny, i in 1:nx
        area_weights[i,j] = cos(deg2rad(latitudes[i,j]))
    end
    
    # Initialize regional sums
    sum_nh = sum_sh = sum_tr = zero(T)
    area_nh = area_sh = area_tr = zero(T)
    
    # Compute regional statistics
    for k in 1:nlevels
        for j in 1:ny
            for i in 1:nx
                weight = area_weights[i,j]
                lat = latitudes[i,j]
                inflation = inflation_factors[i,j,k]
                
                if lat >= config.latbound_nh
                    sum_nh += inflation * weight
                    area_nh += weight
                elseif lat <= config.latbound_sh
                    sum_sh += inflation * weight
                    area_sh += weight
                else
                    sum_tr += inflation * weight
                    area_tr += weight
                end
            end
        end
    end
    
    # Normalize by area
    mean_nh = area_nh > 0 ? sum_nh / area_nh : zero(T)
    mean_sh = area_sh > 0 ? sum_sh / area_sh : zero(T)
    mean_tr = area_tr > 0 ? sum_tr / area_tr : zero(T)
    
    # Global statistics
    min_inflation = minimum(inflation_factors)
    max_inflation = maximum(inflation_factors)
    mean_inflation = mean(inflation_factors)
    
    # Spread statistics
    prior_spread_mean = mean(prior_std)
    posterior_spread_mean = mean(posterior_std)
    final_spread_mean = mean(final_std)
    
    return InflationDiagnostics(
        mean_nh, mean_sh, mean_tr,
        min_inflation, max_inflation, mean_inflation,
        prior_spread_mean, posterior_spread_mean, final_spread_mean,
        area_nh, area_sh, area_tr
    )
end

"""
    adaptive_inflation!(config, innovation_stats, cycle_number)

Adapt inflation parameters based on innovation statistics.
"""
function adaptive_inflation!(
    config::InflationConfiguration{T},
    innovation_stats::Dict{String,T},
    cycle_number::Int
) where T<:AbstractFloat
    
    if !config.adaptive_inflation || cycle_number % config.update_frequency != 0
        return config
    end
    
    # Get innovation statistics
    innovation_mean = get(innovation_stats, "mean", zero(T))
    innovation_std = get(innovation_stats, "std", one(T))
    
    # Adaptive adjustment based on innovation statistics
    adjustment_factor = one(T)
    if abs(innovation_mean) > config.innovation_threshold
        # Increase inflation if innovations are large
        adjustment_factor = T(1.1)
    elseif abs(innovation_mean) < config.innovation_threshold / 2
        # Decrease inflation if innovations are small
        adjustment_factor = T(0.95)
    end
    
    # Update regional weights
    new_analpertwt_nh = clamp(config.analpertwt_nh * adjustment_factor, 0, 2)
    new_analpertwt_sh = clamp(config.analpertwt_sh * adjustment_factor, 0, 2)
    new_analpertwt_tr = clamp(config.analpertwt_tr * adjustment_factor, 0, 2)
    
    return InflationConfiguration(
        new_analpertwt_nh, new_analpertwt_sh, new_analpertwt_tr,
        config.covinflatemin, config.covinflatemax,
        config.smoothparm, config.enable_smoothing,
        config.latbound_nh, config.latbound_sh,
        config.adaptive_inflation, config.innovation_threshold, config.update_frequency,
        config.rtp_inflation, config.rtp_alpha,
        config.save_inflation, config.compute_diagnostics
    )
end

"""
    regional_inflation_weights(latitudes, config)

Convenience function to get regional inflation weights.
"""
function regional_inflation_weights(
    latitudes::AbstractArray{T},
    config::InflationConfiguration{T}
) where T<:AbstractFloat
    return compute_regional_weights(latitudes, config)
end

"""
    smooth_inflation_field(inflation_field, smoothparm, grid_config)

Convenience function for inflation field smoothing.
"""
function smooth_inflation_field(
    inflation_field::Array{T,3},
    smoothparm::T,
    grid_config::GridConfiguration{T}
) where T<:AbstractFloat
    smoothed_field = copy(inflation_field)
    smooth_inflation_field!(smoothed_field, smoothparm, grid_config)
    return smoothed_field
end

"""
    relaxation_to_prior!(ensemble_perturbations, prior_perturbations, alpha)

Apply relaxation-to-prior inflation with specified mixing parameter.
"""
function relaxation_to_prior!(
    ensemble_perturbations::Vector{StateVector{T}},
    prior_perturbations::Vector{StateVector{T}},
    alpha::T
) where T<:AbstractFloat
    
    nensmembers = length(ensemble_perturbations)
    posterior_mean = compute_ensemble_mean(ensemble_perturbations)
    
    # Create temporary config for RTP
    config = InflationConfiguration(
        zero(T), zero(T), zero(T),  # Regional weights (not used in RTP)
        zero(T), T(10),             # Inflation bounds
        zero(T), false,             # Smoothing
        T(20), T(-20),              # Regional boundaries
        false, T(0.1), 1,           # Adaptive parameters
        true, alpha,                # RTP parameters
        false, false                # Diagnostics
    )
    
    # Dummy regional weights (not used in RTP mode)
    nx, ny = 1, 1  # Will be determined from data
    weights = ones(T, nx, ny, 1)
    
    apply_relaxation_to_prior!(ensemble_perturbations, prior_perturbations, 
                              posterior_mean, config, weights)
end

"""
    inflation_diagnostics(result)

Extract inflation diagnostics from InflationResult.
"""
function inflation_diagnostics(result::InflationResult{T}) where T<:AbstractFloat
    return result.diagnostics
end

end # module EnsembleInflation